import torch
import random
import numpy as np
import torch.nn as nn
import torch.optim as optim
from collections import deque

# 优化后的学习参数 - 针对车联网计算卸载场景
LR_ACTOR = 1e-4      # 提高Actor学习率，加快策略收敛
LR_CRITIC = 3e-4     # 保持Critic学习率，但添加学习率衰减
GAMMA = 0.95         # 降低折扣因子，更关注短期奖励（适合车联网动态环境）
MEMORY_SIZE = 2048   # 减小经验池大小，提高样本新鲜度
BATCH_SIZE = 128     # 减小批次大小，提高梯度估计准确性
TAU = 1e-3           # 增大软更新系数，加快目标网络更新

# 新增参数用于提升训练稳定性
NOISE_DECAY = 0.995  # 探索噪声衰减因子
MIN_NOISE = 0.01     # 最小探索噪声
GRAD_CLIP = 0.5      # 梯度裁剪阈值

HIDDEN_DIM = 128

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# print(device)

class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=HIDDEN_DIM):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, hidden_dim)
        self.fc5 = nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = torch.softmax(self.fc5(x), dim=-1)
        return x


class Critic(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=HIDDEN_DIM):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, 1)

    def forward(self, x, a):
        x = torch.cat([x, a], 1)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        return self.fc4(x)


class ReplayMemory:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def add_memo(self, state, action, reward, next_state):
        state = np.expand_dims(state, 0)
        next_state = np.expand_dims(next_state, 0)
        self.buffer.append((state, action, reward, next_state))

    def sample(self, batch_size):
        state, action, reward, next_state = zip(*random.sample(self.buffer, batch_size))
        return np.concatenate(state), action, reward, np.concatenate(next_state)

    def __len__(self):
        return len(self.buffer)


class DDPGAgent:
    def __init__(self, state_dim, action_dim):
        self.actor = Actor(state_dim, action_dim).to(device)
        self.actor_target = Actor(state_dim, action_dim).to(device)
        self.actor_target.load_state_dict(self.actor.state_dict())
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=LR_ACTOR)

        self.critic = Critic(state_dim, action_dim).to(device)
        self.critic_target = Critic(state_dim, action_dim).to(device)
        self.critic_target.load_state_dict(self.critic.state_dict())
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=LR_CRITIC)

        # 添加学习率调度器 - 针对车联网动态环境优化
        self.actor_scheduler = optim.lr_scheduler.ExponentialLR(self.actor_optimizer, gamma=0.99)
        self.critic_scheduler = optim.lr_scheduler.ExponentialLR(self.critic_optimizer, gamma=0.995)
        
        self.replay_buffer = ReplayMemory(MEMORY_SIZE)
        
        # 探索噪声参数
        self.noise_scale = 0.1
        self.update_count = 0

    def get_action(self, state, add_noise=True):
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        with torch.no_grad():
            action_probs = self.actor(state)
            
        if add_noise and self.noise_scale > MIN_NOISE:
            # 添加自适应探索噪声 - 基于当前动作概率的不确定性
            entropy = -torch.sum(action_probs * torch.log(action_probs + 1e-8))
            # 熵越小（越确定），添加更多噪声；熵越大（越不确定），添加较少噪声
            adaptive_noise = self.noise_scale * (1.0 + 0.5 * (1.0 - entropy.item()))
            
            # 使用温度参数调整探索
            temperature = max(0.1, 1.0 - self.update_count / 10000.0)
            action_probs = torch.softmax(torch.log(action_probs + 1e-8) / temperature, dim=-1)
            
            # 添加Dirichlet噪声以增强探索
            noise = torch.distributions.Dirichlet(torch.ones_like(action_probs) * adaptive_noise).sample()
            action_probs = 0.75 * action_probs + 0.25 * noise
            
        dist = torch.distributions.Categorical(action_probs)
        action_discrete_index = dist.sample()
        action_discrete = [0.0 for _ in range(6)]
        action_discrete[action_discrete_index.item()] = 1.0
        return action_discrete

    def update(self):
        if len(self.replay_buffer) < BATCH_SIZE:
            return

        states, actions, rewards, next_states = self.replay_buffer.sample(BATCH_SIZE)
        states = torch.FloatTensor(states).to(device)
        actions = torch.FloatTensor(np.vstack(actions)).to(device)
        rewards = torch.FloatTensor(rewards).unsqueeze(1).to(device)
        next_states = torch.FloatTensor(next_states).to(device)

        # 奖励标准化 - 提升训练稳定性
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-8)

        # update critic with improved stability
        with torch.no_grad():
            next_actions = self.actor_target(next_states)
            target_Q = self.critic_target(next_states, next_actions)
            target_Q = rewards + (GAMMA * target_Q)
            
        current_Q = self.critic(states, actions)
        critic_loss = nn.MSELoss()(current_Q, target_Q)
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        # 梯度裁剪 - 防止梯度爆炸
        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), GRAD_CLIP)
        self.critic_optimizer.step()

        # update actor with delayed policy updates (每2次更新critic才更新一次actor)
        if self.update_count % 2 == 0:
            actor_loss = -self.critic(states, self.actor(states)).mean()
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(self.actor.parameters(), GRAD_CLIP)
            self.actor_optimizer.step()

            # update target networks with improved soft update
            for target_param, param in zip(self.actor_target.parameters(), self.actor.parameters()):
                target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)
            for target_param, param in zip(self.critic_target.parameters(), self.critic.parameters()):
                target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)

        self.update_count += 1
        
        # 学习率调度 - 每100次更新调整一次
        if self.update_count % 100 == 0:
            self.actor_scheduler.step()
            self.critic_scheduler.step()
            
        # 探索噪声衰减
        self.noise_scale = max(MIN_NOISE, self.noise_scale * NOISE_DECAY)
